#include "container.hpp"
#include <cmath>
#include <stdexcept>

Mesh::Mesh(int nx, int ny, double dx, double dy) 
    : nx_(nx), ny_(ny), dx_(dx), dy_(dy), 
      solution_((nx+1)*(ny+1), 0.0) {
		rhs_ = new double[(nx+1)*(ny+1)];
	  }

void Mesh::setBoundaryConditions(double top, double bottom, double left, double right) {
    // Set boundary values
    for (int i = 0; i <= nx_; ++i) {
        solution_[i] = bottom; // bottom boundary
        solution_[i + ny_*(nx_+1)] = top; // top boundary
    }
    for (int j = 0; j <= ny_; ++j) {
        solution_[j*(nx_+1)] = left; // left boundary
        solution_[nx_ + j*(nx_+1)] = right; // right boundary
    }
}

void Mesh::setBoundaryCondition(int i, int j, double value) {
    if (i < 0 || i > nx_ || j < 0 || j > ny_) {
        throw std::out_of_range("Boundary condition index out of range");
    }
    int idx = i + j*(nx_+1);
    solution_[idx] = value;
}

void Mesh::solvePoisson(int max_iterations, double tolerance) {
    for (int iter = 0; iter < max_iterations; ++iter) {
        std::vector<double> old_solution = solution_;
        applyStencil();
        
        // Check convergence
        double max_diff = 0.0;
        for (size_t i = 0; i < solution_.size(); ++i) {
            max_diff = std::max(max_diff, std::abs(solution_[i] - old_solution[i]));
        }
        
        if (max_diff < tolerance) {
            break;
        }
    }
}

double Mesh::getValue(int i, int j) const {
    if (i < 0 || i > nx_ || j < 0 || j > ny_) {
        throw std::out_of_range("Index out of range");
    }
    int idx = i + j*(nx_+1);
    return solution_[idx];
}

void Mesh::applyStencil() {
    // Simple 5-point stencil for Poisson equation
    std::vector<double> new_solution = solution_;
    for (int j = 1; j < ny_; ++j) {
        for (int i = 1; i < nx_; ++i) {
            int idx = i + j*(nx_+1);
            new_solution[idx] = 0.25 * (
                solution_[idx-1] + solution_[idx+1] + 
                solution_[idx-(nx_+1)] + solution_[idx+(nx_+1)] - 
                dx_*dy_*rhs_[idx]
            );
        }
    }
    solution_ = new_solution;
}

const std::vector<double>& Mesh::getSolution() const {
    return solution_;
}


Mesh::~Mesh()
{
	delete[] rhs_;
	rhs_ = nullptr;
}
